from jose import jwt
from fastapi.security import OAuth2PasswordBearer
from fastapi import HTTPException, status, Request, Header
from ..model.slave import User
from ..utils.database import SessionLocal
from typing import Union
from datetime import datetime, timedelta
    

class MyOAuth2PasswordBearer(OAuth2PasswordBearer):
    def __init__(self, tokenUrl: str):
        super().__init__(
            tokenUrl=tokenUrl,
            scheme_name="",
            scopes=None,
            description="JWT Authentication",
            auto_error=True
        )
        self.SECRET_KEY = "ed970259a19edfedf1010199c7002d183bd15bcaec612481b29bac1cb83d8137"
        self.ALGORITHM = "HS256"
        self.ACCESS_TOKEN_EXPIRE_MINUTES = 30

        self.credentials_exception = HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Could not validate credentials",
            headers={"WWW-Authenticate": "Bearer"},
        )
         
    def create_jwt_token(self, data: dict, expires_delta: Union[timedelta, None] = None):
        # 设置 token 过期时间
        if expires_delta:
            expire = datetime.utcnow() + expires_delta
            data.update({"exp": expire})
        
        # 进行jwt加密
        token = jwt.encode(claims=data, key=self.SECRET_KEY, algorithm=self.ALGORITHM)
        return token

    def check_jwt_token(self, token):
        try:
            pyload = jwt.decode(
                token,
                self.SECRET_KEY,
                algorithms=[self.ALGORITHM]
            )

            user_name: str = pyload.get("user_name")

            with SessionLocal() as session:
                self.user = session.query(User).filter_by(user_name=user_name).first()

            if not self.user:
                raise self.credentials_exception
            return self.user
        except (jwt.JWTError, jwt.ExpiredSignatureError, AttributeError):
            # 抛出自定义异常， 然后捕获统一响应
            raise self.credentials_exception 
         
    async def __call__(self, request: Request):
        path: str = request.url.path

         # 定义无需认证的路径
        excluded_paths = {
            "/login",
            "/register",
            "/docs",
            "/"
        }

        # 如果路径在豁免列表中，跳过认证
        if any(path.endswith(excluded) for excluded in excluded_paths):
            return ""
        
        try:
            authorization = request.headers["Authorization"]

            scheme, _, token = authorization.partition(" ")

            if scheme.lower() != "bearer":
                token = authorization

        except Exception:
            raise self.credentials_exception
        
        self.user = self.check_jwt_token(token)

        return self.user
    

my_oauth2_scheme = MyOAuth2PasswordBearer(tokenUrl='api/user/login')